package com.gitee.huanminabc.jcommon.multithreading;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;

/**
 * 同一个 key 可以有 n个线程执行,超过n个线程则阻塞, 意思就是对key进行限流
 * @Author huanmin
 * @Date 2024/6/5 下午1:08
 */
public class SemaphoreLockByKey implements LockKey<String> {

    private final Map<String, Semaphore> semaphores = new ConcurrentHashMap<>();

    /**
     * 最大线程,意思是同一个key可以有多少个线程同时执行
     */
    private final int allowed_threads;

    public SemaphoreLockByKey(int allowed_threads) {
        this.allowed_threads = allowed_threads;
    }

    /**
     * 加锁
     */
    @Override
    public void lock(String key) {
        Semaphore semaphore = semaphores.computeIfAbsent(key, k -> new Semaphore(allowed_threads));
        semaphore.acquireUninterruptibly();//获取许可
    }

    @Override
    public boolean tryLock(String key) {
        Semaphore semaphore1 = semaphores.computeIfAbsent(key ,k->new Semaphore(allowed_threads));
        return semaphore1.tryAcquire();
    }

    @Override
    public boolean tryLock(String key, long timeout, TimeUnit unit) {
        Semaphore semaphore1 = semaphores.computeIfAbsent(key, k -> new Semaphore(allowed_threads));
        try {
            return semaphore1.tryAcquire(timeout, unit);
        } catch (InterruptedException e) {
            return false;
        }
    }


    /**
     * 解锁
     */
    @Override
    public void unlock(String key) {
        // 如果key为空，直接返回
        if (key == null) {
            throw new IllegalArgumentException("key 不能为空");
        }

        // 从Map中获取锁对象
        Semaphore semaphore = semaphores.get(key);
        if (semaphore == null) {
            throw new IllegalArgumentException("key " + key + "尚未加锁");
        }
        //归还许可
        semaphore.release();

    }
}